import os

def filter_and_remap_wheel_labels(label_dir, output_dir=None):
    """
    从YOLO标签中只保留 wheel (原class_id=2)，并将其重映射为 class_id=0
    :param label_dir: 原始标签文件夹路径
    :param output_dir: 输出路径，None 表示覆盖原文件
    """
    if output_dir is None:
        output_dir = label_dir
        os.makedirs(output_dir, exist_ok=True)
    else:
        os.makedirs(output_dir, exist_ok=True)

    for filename in os.listdir(label_dir):
        if not filename.endswith('.txt'):
            continue

        label_path = os.path.join(label_dir, filename)
        output_path = os.path.join(output_dir, filename)

        with open(label_path, 'r') as f:
            lines = f.readlines()

        # 保留原 class_id == 2 的行，并将其改为 0
        remapped_lines = []
        for line in lines:
            parts = line.strip().split()
            if len(parts) < 5:
                continue  # 跳过格式错误的行
            class_id = parts[0]
            if class_id == '2':
                # 修改类别为 0，保留其余坐标
                new_line = '0 ' + ' '.join(parts[1:]) + '\n'
                remapped_lines.append(new_line)

        # 写入输出文件
        with open(output_path, 'w') as f:
            f.writelines(remapped_lines)

        print(f"Processed: {filename} -> 保留 {len(remapped_lines)} 个 wheel 标注（已重映射为类别 0）")

# 使用示例
if __name__ == "__main__":
    label_directory = "D:\Datasets\dataset2CheckNew\labels-原始-三类别"          # 原始标签路径
    output_directory = "D:\Datasets\dataset2CheckNew\labels-一类别wheel"   # 新标签路径（建议不要覆盖）

    filter_and_remap_wheel_labels(label_directory, output_directory)